package com.example.springai.controller;

import com.example.springai.bean.ChatTypeEnum;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.model.Media;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.util.CollectionUtils;
import org.springframework.util.MimeType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Objects;
import java.util.function.Consumer;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@RestController
@RequestMapping("/ai")
@Slf4j
public class ChatController {

    @Autowired
    private ChatClient chatClient;

    @Autowired
    @Qualifier("inDatabaseChatHistoryRepository")
    private com.example.springai.repository.ChatHistoryRepository ChatHistoryRepository;

    /**
     * 流式返回结果
     *
     * @param prompt 用户输入的提示词
     * @return 大模型处理结果
     */
    @PostMapping(value = "/chat", produces = "text/html;charset=utf-8")
    public Flux<String> chatByClientStream(@RequestParam String prompt, @RequestParam String chatId, @RequestParam(required = false, value = "files") List<MultipartFile> files) {
        log.info("用户输入: {}", prompt);

        // [1] 保存会话id到数据库
        ChatHistoryRepository.save(ChatTypeEnum.CHAT.getValue(), chatId);

        if (CollectionUtils.isEmpty(files)) {
            return textChat(prompt, chatId);
        } else {
            return fileChat(prompt, chatId, files);
        }
    }


    // 处理文件聊天
    public Flux<String> fileChat(String prompt, String chatId, List<MultipartFile> files) {

        // 1.解析多媒体
        List<Media> medias = files.stream()
                .map(file -> new Media(
                                MimeType.valueOf(Objects.requireNonNull(file.getContentType())),
                                file.getResource()
                        )
                )
                .toList();


        return chatClient.prompt()
                .user(promptUserSpec -> promptUserSpec.text(prompt).media(medias.toArray(Media[]::new)))
                .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }


    // 文本聊天
    public Flux<String> textChat(String prompt, String chatId) {
        return chatClient.prompt()
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }
}
